import torch
import cv2
import os
import torch.nn as nn
import numpy as np
from .model import Generator
import matplotlib.pyplot as plt
import matplotlib.cm as cm


def weights_init(net):
    for m in net.modules():
        if isinstance(m, nn.Conv2d):
            m.weight.data.normal_(0.0, 0.01)
            if m.bias is not None:
                m.bias.data.fill_(0.0)
        elif isinstance(m, nn.Linear):
            m.weight.data.normal_(0.0, 0.01)
            if m.bias is not None:
                m.bias.data.fill_(0.0)
        elif isinstance(m, nn.BatchNorm2d):
            m.weight.data.normal_(1.0, 0.01)
            if m.bias is not None:
                m.bias.data.fill_(0.0)


def evaluate_model(trained_model, data_loader, is_cuda):
    net = Generator()
    net.load_state_dict(torch.load(trained_model))
    if is_cuda:
        net.cuda()
    net.eval()
    mae = 0.0
    mse = 0.0
    for blob in data_loader:
        im_data = blob['img_data']
        gt_data = blob['gt_data']
        if is_cuda:
            im_data = im_data.cuda()
        density_map = net(im_data)
        density_map = density_map.data.cpu().numpy()
        gt_data = (gt_data.data.numpy() + 1.0) / 2.0
        density_map = (density_map + 1.0) / 2.0
        gt_count = np.sum(gt_data)
        et_count = np.sum(density_map)
        mae += abs(gt_count-et_count)
        mse += ((gt_count-et_count)*(gt_count-et_count))

    return mae, mse


def save_results(input_img, gt_data, density_map, output_dir, fname='results.png'):
    input_img = input_img[0][0]
    gt_data = 255*gt_data/np.max(gt_data)
    density_map = 255*density_map/np.max(density_map)
    gt_data = gt_data[0][0]
    density_map = density_map[0][0]
    if density_map.shape[1] != input_img.shape[1]:
        density_map = cv2.resize(density_map, (input_img.shape[1], input_img.shape[0]))
        gt_data = cv2.resize(gt_data, (input_img.shape[1], input_img.shape[0]))
    result_img = np.hstack((input_img, gt_data, density_map))
    cv2.imwrite(os.path.join(output_dir, fname), result_img)


def collate_fn(batch):
    return batch[0]


def save_density_map(density_map, output_dir, fname='results.png', et_count=0):
    density_map = 255*density_map/np.max(density_map)
    density_map = density_map[0][0]

    plt.figure()
    plt.imshow(density_map, cmap=cm.jet)
    plt.colorbar()
    plt.title('sum:'+str(et_count))
    plt.savefig(os.path.join(output_dir, fname))


def show_loss_acc(epoch_x, loss_list_d, loss_list_g, file_name=None):
    plt.figure()
    plt.plot(epoch_x, loss_list_d, label='loss_D', color='r')
    plt.plot(epoch_x, loss_list_g, label='loss_G', color='g')
    plt.xlabel('epoch')
    plt.legend()
    if file_name is None:
        plt.show()
    else:
        plt.savefig(file_name)
